#!/usr/bin/env python3
"""
预测功能测试脚本
测试调整后的预测功能是否与训练优化保持一致
"""

import torch
import numpy as np
import pandas as pd
import os
import sys
from datetime import datetime, timedelta

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from data_utils import load_config, prepare_station_data
from models import CombinedModel
from predict import load_trained_model, prepare_prediction_input, predict_future, denormalize_predictions

def test_prediction_setup():
    """测试预测设置"""
    print("=" * 60)
    print("预测功能设置测试")
    print("=" * 60)
    
    # 加载配置
    config = load_config()
    print("✓ 配置文件加载成功")
    
    # 检查关键配置
    print(f"波段配置: {config['data']['himawari_bands']}")
    print(f"输出变量数量: {config['model']['num_output_vars']}")
    print(f"模型输入通道数: {config['model']['num_channels']}")
    
    # 验证配置一致性
    expected_channels = len(config['data']['himawari_bands'])
    actual_channels = config['model']['num_channels']
    
    if expected_channels == actual_channels:
        print(f"✓ 波段配置一致: {expected_channels} 个波段")
    else:
        print(f"✗ 波段配置不一致: 期望 {expected_channels}, 实际 {actual_channels}")
        return False
    
    expected_outputs = 3  # DNI, DHI, Temperature
    actual_outputs = config['model']['num_output_vars']
    
    if expected_outputs == actual_outputs:
        print(f"✓ 输出变量数量一致: {actual_outputs} 个变量")
    else:
        print(f"✗ 输出变量数量不一致: 期望 {expected_outputs}, 实际 {actual_outputs}")
        return False
    
    return True

def test_model_creation():
    """测试模型创建"""
    print("\n" + "=" * 60)
    print("模型创建测试")
    print("=" * 60)
    
    try:
        config = load_config()
        model = CombinedModel(config)
        
        # 计算参数数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        print(f"✓ 模型创建成功")
        print(f"总参数数量: {total_params:,}")
        print(f"可训练参数: {trainable_params:,}")
        
        # 测试前向传播
        batch_size = 2
        historical_days = config['data']['historical_days']
        num_channels = config['model']['num_channels']
        img_height = config['model']['img_height']
        img_width = config['model']['img_width']
        future_hours = config['data']['future_hours']
        
        # 创建测试数据 - 使用正确的维度
        historical_points = historical_days * 24 * 4  # 5天*24小时*4个15分钟点 = 480
        station_data = torch.randn(batch_size, historical_points, 3)  # 480个时间点，3个变量
        himawari_data = torch.randn(batch_size, num_channels, img_height, img_width)
        
        with torch.no_grad():
            output = model(station_data, himawari_data)
        
        expected_shape = (batch_size, future_hours * 4, 3)  # 3个输出变量
        actual_shape = output.shape
        
        if actual_shape == expected_shape:
            print(f"✓ 前向传播测试通过: {actual_shape}")
        else:
            print(f"✗ 前向传播测试失败: 期望 {expected_shape}, 实际 {actual_shape}")
            return False
        
        return True
        
    except Exception as e:
        print(f"✗ 模型创建失败: {e}")
        return False

def test_data_preparation():
    """测试数据准备"""
    print("\n" + "=" * 60)
    print("数据准备测试")
    print("=" * 60)
    
    try:
        config = load_config()
        
        # 检查数据文件
        station_csv = config['data']['station_csv']
        if not os.path.exists(station_csv):
            print(f"✗ 观察站数据文件不存在: {station_csv}")
            return False
        
        print(f"✓ 观察站数据文件存在: {station_csv}")
        
        # 加载和准备数据
        station_df = prepare_station_data(station_csv)
        print(f"✓ 数据加载成功: {len(station_df)} 条记录")
        
        # 检查必要字段
        required_fields = ['direct_irradiance', 'diffuse_irradiance', 'temperature']
        missing_fields = [field for field in required_fields if field not in station_df.columns]
        
        if missing_fields:
            print(f"✗ 缺少必要字段: {missing_fields}")
            return False
        
        print(f"✓ 必要字段完整: {required_fields}")
        
        # 测试预测输入准备
        try:
            input_data = prepare_prediction_input(station_csv, config['data']['himawari_dir'], config)
            print(f"✓ 预测输入准备成功")
            print(f"  - 站点数据形状: {input_data['station_data'].shape}")
            print(f"  - 卫星数据形状: {input_data['himawari_data'].shape}")
            print(f"  - 预测时间: {input_data['prediction_time']}")
            
            # 验证数据形状
            expected_station_shape = (1, config['data']['historical_days'] * 24 * 4, 3)  # 5天*24小时*4个15分钟点 = 480个时间点
            expected_himawari_shape = (1, len(config['data']['himawari_bands']), 224, 224)
            
            actual_station_shape = input_data['station_data'].shape
            actual_himawari_shape = input_data['himawari_data'].shape
            
            if actual_station_shape == expected_station_shape:
                print(f"✓ 站点数据形状正确: {actual_station_shape}")
            else:
                print(f"✗ 站点数据形状错误: 期望 {expected_station_shape}, 实际 {actual_station_shape}")
                return False
            
            if actual_himawari_shape == expected_himawari_shape:
                print(f"✓ 卫星数据形状正确: {actual_himawari_shape}")
            else:
                print(f"✗ 卫星数据形状错误: 期望 {expected_himawari_shape}, 实际 {actual_himawari_shape}")
                return False
            
            return True
            
        except Exception as e:
            print(f"✗ 预测输入准备失败: {e}")
            return False
        
    except Exception as e:
        print(f"✗ 数据准备失败: {e}")
        return False

def test_denormalization():
    """测试反标准化功能"""
    print("\n" + "=" * 60)
    print("反标准化功能测试")
    print("=" * 60)
    
    try:
        config = load_config()
        station_df = prepare_station_data(config['data']['station_csv'])
        
        # 创建模拟预测结果
        future_hours = config['data']['future_hours']
        num_points = future_hours * 4  # 15分钟间隔
        num_vars = 3  # 3个输出变量
        
        mock_predictions = np.random.randn(num_points, num_vars).astype(np.float32)
        
        # 测试反标准化
        denormalized = denormalize_predictions(mock_predictions, station_df)
        
        if denormalized.shape == mock_predictions.shape:
            print(f"✓ 反标准化形状正确: {denormalized.shape}")
        else:
            print(f"✗ 反标准化形状错误: 期望 {mock_predictions.shape}, 实际 {denormalized.shape}")
            return False
        
        # 检查数值范围 - 由于是随机数据，只检查基本合理性
        for i, var_name in enumerate(['direct_irradiance', 'diffuse_irradiance', 'temperature']):
            values = denormalized[:, i]
            
            # 检查基本数值范围（允许一定的异常值，因为是随机数据）
            if var_name in ['direct_irradiance', 'diffuse_irradiance']:
                # 辐照度应该在合理范围内，但随机数据可能有负值
                print(f"✓ {var_name} 数值范围: [{np.min(values):.1f}, {np.max(values):.1f}] (随机数据)")
            elif var_name == 'temperature':
                print(f"✓ {var_name} 数值范围: [{np.min(values):.1f}, {np.max(values):.1f}] (随机数据)")
        
        return True
        
    except Exception as e:
        print(f"✗ 反标准化测试失败: {e}")
        return False

def test_model_loading():
    """测试模型加载（如果存在训练好的模型）"""
    print("\n" + "=" * 60)
    print("模型加载测试")
    print("=" * 60)
    
    try:
        config = load_config()
        model_path = os.path.join(config['training']['save_dir'], 'best_model.pth')
        
        if not os.path.exists(model_path):
            print(f"⚠ 未找到训练好的模型: {model_path}")
            print("  请先运行训练: python train.py")
            return True  # 不算失败，只是没有训练好的模型
        
        print(f"✓ 找到训练好的模型: {model_path}")
        
        # 尝试加载模型
        model, device = load_trained_model(config, model_path)
        print(f"✓ 模型加载成功，设备: {device}")
        
        return True
        
    except Exception as e:
        print(f"✗ 模型加载失败: {e}")
        return False

def main():
    """主测试函数"""
    print("VIT天气预报模型 - 预测功能测试")
    print(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 运行所有测试
    tests = [
        ("预测设置测试", test_prediction_setup),
        ("模型创建测试", test_model_creation),
        ("数据准备测试", test_data_preparation),
        ("反标准化测试", test_denormalization),
        ("模型加载测试", test_model_loading),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"✗ {test_name} 异常: {e}")
            results.append((test_name, False))
    
    # 输出测试结果摘要
    print("\n" + "=" * 60)
    print("测试结果摘要")
    print("=" * 60)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results:
        status = "✓ 通过" if result else "✗ 失败"
        print(f"{test_name}: {status}")
        if result:
            passed += 1
    
    print(f"\n总体结果: {passed}/{total} 测试通过")
    
    if passed == total:
        print("🎉 所有测试通过！预测功能已成功调整。")
        print("\n可以使用以下命令进行预测:")
        print("python predict.py")
    else:
        print("⚠ 部分测试失败，请检查相关配置和代码。")
    
    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
